package com.starry.module.system.core.oauth2.authorization.grant;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.json.JSONUtil;
import com.starry.core.common.constants.CommonConstant;
import com.starry.module.system.core.oauth2.authorization.constant.SecurityConstants;
import com.starry.module.system.core.oauth2.authorization.third.ThireClient;
import com.starry.module.system.core.oauth2.authorization.third.ThireClientFactory;
import com.starry.module.system.core.oauth2.authorization.utils.SecurityUtils;
import com.starry.module.system.core.oauth2.entity.Oauth2ThirdAccount;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.OidcScopes;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;
import org.springframework.util.ObjectUtils;

import java.security.Principal;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 自定义Grant登录认证提供者
 *
 * @author 小柯
 */
@Slf4j
public class CustomGrantAuthenticationProvider implements AuthenticationProvider {

    private static final String ERROR_URI = "https://datatracker.ietf.org/doc/html/rfc6749#section-5.2";
    private static final OAuth2TokenType ID_TOKEN_TOKEN_TYPE = new OAuth2TokenType(OidcParameterNames.ID_TOKEN);
    private OAuth2TokenGenerator<?> tokenGenerator;
    private AuthenticationManager authenticationManager;
    private OAuth2AuthorizationService authorizationService;
    private final ThireClientFactory thireClientFactory;

    public CustomGrantAuthenticationProvider(ThireClientFactory thireClientFactory) {
        this.thireClientFactory = thireClientFactory;
    }


    @Override
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {
        CustomGrantAuthenticationToken authenticationToken = (CustomGrantAuthenticationToken) authentication;

        //确保客户端经过身份验证
        OAuth2ClientAuthenticationToken clientPrincipal =
                SecurityUtils.getAuthenticatedClientElseThrowInvalidClient(authenticationToken);
        RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();
        // 确保客户端配置为使用此授权授予类型
        if (registeredClient == null || !registeredClient.getAuthorizationGrantTypes().contains(authenticationToken.getAuthorizationGrantType())) {
            throw new OAuth2AuthenticationException(OAuth2ErrorCodes.UNAUTHORIZED_CLIENT);
        }


        // 验证scope
        Set<String> authorizedScopes = getAuthorizedScopes(registeredClient, authenticationToken.getScopes());

        // 进行认证
        Authentication authenticate = getAuthenticatedUser(authenticationToken, authenticationToken.getAuthorizationGrantType());

        // 以下内容摘抄自OAuth2AuthorizationCodeAuthenticationProvider
        DefaultOAuth2TokenContext.Builder tokenContextBuilder = DefaultOAuth2TokenContext.builder()
                .registeredClient(registeredClient)
                .principal(authenticate)
                .authorizationServerContext(AuthorizationServerContextHolder.getContext())
                .authorizedScopes(authorizedScopes)
                .authorizationGrantType(authenticationToken.getAuthorizationGrantType())
                .authorizationGrant(authenticationToken);


        // Initialize the OAuth2Authorization
        OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
                // 存入授权scope
                .authorizedScopes(authorizedScopes)
                // 当前授权用户名称
                .principalName(authenticate.getName())
                // 设置当前用户认证信息
                .attribute(Principal.class.getName(), authenticate)
                .authorizationGrantType(authenticationToken.getAuthorizationGrantType());

        // ----- Access token -----
        OAuth2TokenContext tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.ACCESS_TOKEN).build();
        OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
        if (generatedAccessToken == null) {
            OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
                    "The token generator failed to generate the access token.", ERROR_URI);
            throw new OAuth2AuthenticationException(error);
        }


        // ----- Access Token -----
        OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
                generatedAccessToken.getTokenValue(), generatedAccessToken.getIssuedAt(),
                generatedAccessToken.getExpiresAt(), tokenContext.getAuthorizedScopes());
        if (generatedAccessToken instanceof ClaimAccessor) {
            authorizationBuilder.token(accessToken, (metadata) ->
                    metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, ((ClaimAccessor) generatedAccessToken).getClaims()));
        } else {
            authorizationBuilder.accessToken(accessToken);
        }

        // ----- Refresh token -----
        OAuth2RefreshToken refreshToken = null;
        if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN) &&
                !clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {

            tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build();
            OAuth2Token generatedRefreshToken = this.tokenGenerator.generate(tokenContext);
            if (!(generatedRefreshToken instanceof OAuth2RefreshToken)) {
                OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
                        "令牌生成器无法生成刷新令牌", ERROR_URI);
                throw new OAuth2AuthenticationException(error);
            }
            refreshToken = (OAuth2RefreshToken) generatedRefreshToken;
            authorizationBuilder.refreshToken(refreshToken);
        }

        // ----- ID token -----
        OidcIdToken idToken;
        if (authorizedScopes.contains(OidcScopes.OPENID)) {
            tokenContext = tokenContextBuilder
                    .tokenType(ID_TOKEN_TOKEN_TYPE)
                    .authorization(authorizationBuilder.build())
                    .build();

            OAuth2Token generatedIdToken = this.tokenGenerator.generate(tokenContext);
            if (!(generatedIdToken instanceof Jwt)) {
                OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
                        "The token generator failed to generate the ID token.", ERROR_URI);
                throw new OAuth2AuthenticationException(error);
            }

            idToken = new OidcIdToken(generatedIdToken.getTokenValue(), generatedIdToken.getIssuedAt(),
                    generatedIdToken.getExpiresAt(), ((Jwt) generatedIdToken).getClaims());
            authorizationBuilder.token(idToken, (metadata) ->
                    metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, idToken.getClaims()));
        } else {
            idToken = null;
        }
        OAuth2Authorization authorization = authorizationBuilder.build();


        // 保存到数据库中
        this.authorizationService.save(authorization);

        Map<String, Object> additionalSchemas = new HashMap<>(1);
        if (idToken != null) {
            // 放入idToken
            additionalSchemas.put(OidcParameterNames.ID_TOKEN, idToken.getTokenValue());
        }

        return new OAuth2AccessTokenAuthenticationToken(registeredClient, clientPrincipal, accessToken, refreshToken, additionalSchemas);
    }

    /**
     * 获取认证过的scope
     *
     * @param registeredClient 客户端
     * @param requestedScopes  请求中的scope
     * @return 认证过的scope
     */
    private Set<String> getAuthorizedScopes(RegisteredClient registeredClient, Set<String> requestedScopes) {
        // Default to configured scopes
        Set<String> authorizedScopes = registeredClient.getScopes();
        if (!ObjectUtils.isEmpty(requestedScopes)) {
            Set<String> unauthorizedScopes = requestedScopes.stream()
                    .filter(requestedScope -> !registeredClient.getScopes().contains(requestedScope))
                    .collect(Collectors.toSet());
            if (!ObjectUtils.isEmpty(unauthorizedScopes)) {
                SecurityUtils.throwError(
                        OAuth2ErrorCodes.INVALID_REQUEST,
                        "OAuth 2.0 Schema: " + OAuth2ParameterNames.SCOPE,
                        ERROR_URI);
            }
            authorizedScopes = new LinkedHashSet<>(requestedScopes);
        }
        return authorizedScopes;
    }

    /**
     * 获取认证过的用户信息
     *
     * @param authenticationToken converter构建的认证信息，这里是包含手机号与验证码的
     * @return 认证信息
     */
    public Authentication getAuthenticatedUser(CustomGrantAuthenticationToken authenticationToken, AuthorizationGrantType authorizationGrantType) {
        // 内部认证客户端模式
        // 获取手机号密码
        Map<String, Object> additionalSchemas = authenticationToken.getAdditionalSchemas();

        // 构建UsernamePasswordAuthenticationToken通过AbstractUserDetailsAuthenticationProvider及其子类对手机号与验证码进行校验
        // 这里就是我说的短信验证与密码模式区别不大，如果是短信验证模式则在SmsCaptchaLoginAuthenticationProvider中加一个校验，
        // 使框架支持手机号、验证码校验，反之不加就是账号密码登录
        String grantType = authorizationGrantType.getValue();
        UsernamePasswordAuthenticationToken unauthenticated = null;
        // 内部认证客户端模式
        if (grantType.equals(CommonConstant.INTERNAL_CLIENT_GRANT_TYPE)) {
            unauthenticated = UsernamePasswordAuthenticationToken.unauthenticated(CommonConstant.INTERNAL_CLIENT_CLIENT_SECRET, "");
        }
        // 短信验证模式
        if (grantType.equals(SecurityConstants.GRANT_TYPE_SMS_CODE)) {
            String phone = (String) additionalSchemas.get(SecurityConstants.OAUTH_SCHEMA_NAME_PHONE);
            String smsCaptcha = (String) additionalSchemas.get(SecurityConstants.OAUTH_Schema_NAME_SMS_CAPTCHA);
            //校验验证码
            Oauth2ThirdAccount oauth2ThirdAccount = new Oauth2ThirdAccount();
            oauth2ThirdAccount.setOpenid(phone);
            unauthenticated = UsernamePasswordAuthenticationToken.unauthenticated(JSONUtil.toJsonStr(oauth2ThirdAccount), smsCaptcha);
        }
        // 密码验证模式
        if (grantType.equals(SecurityConstants.GRANT_TYPE_PASSWORD)) {
            String username = (String) additionalSchemas.get("username");
            String password = (String) additionalSchemas.get("password");

            Oauth2ThirdAccount oauth2ThirdAccount = new Oauth2ThirdAccount();
            oauth2ThirdAccount.setOpenid(username);

            unauthenticated = UsernamePasswordAuthenticationToken.unauthenticated(JSONUtil.toJsonStr(oauth2ThirdAccount), password);
        }

        // 第三方验证模式
        if (grantType.equals(SecurityConstants.GRANT_TYPE_THIRD)) {
            String code = (String) additionalSchemas.get("code");
            if (!ObjectUtil.isAllNotEmpty(code)) {
                SecurityUtils.throwError(
                        OAuth2ErrorCodes.INVALID_REQUEST,
                        "第三方授权码不能为空",
                        ERROR_URI
                );
            }

            String state = (String) additionalSchemas.get("state");

            String type = (String) additionalSchemas.get("type");
            if (!ObjectUtil.isAllNotEmpty(type)) {
                SecurityUtils.throwError(
                        OAuth2ErrorCodes.INVALID_REQUEST,
                        "第三方登录类型不能为空",
                        ERROR_URI
                );
            }

            ThireClient thireClient = null;
            try {
                thireClient = thireClientFactory.getClient(type);
            } catch (Exception e) {
                log.error("获取第三方认证客户端失败", e);
                // 认证失败
                SecurityUtils.throwError(
                        OAuth2ErrorCodes.INVALID_REQUEST,
                        e.getMessage(),
                        ERROR_URI
                );
            }

            assert thireClient != null;

            // 获取第三方用户信息
            try {
                Oauth2ThirdAccount oauth2ThirdAccount = thireClient.loadUser(code, state);
                unauthenticated = UsernamePasswordAuthenticationToken.unauthenticated(JSONUtil.toJsonStr(oauth2ThirdAccount), "");
            } catch (Exception e) {
                log.error("第三方认证失败", e);
                SecurityUtils.throwError(
                        OAuth2ErrorCodes.INVALID_REQUEST,
                        e.getMessage(),
                        ERROR_URI
                );
            }
        }

        Authentication authenticate = null;
        try {
            authenticate = authenticationManager.authenticate(unauthenticated);
        } catch (Exception e) {
            log.error("认证失败", e);
            // 认证失败
            String message = e.getMessage().equals("Bad credentials") ? "账号或密码错误" : e.getMessage();
            SecurityUtils.throwError(
                    OAuth2ErrorCodes.INVALID_REQUEST,
                    message,
                    ERROR_URI
            );
        }
        return authenticate;
    }

    @Override
    public boolean supports(Class<?> authentication) {
        return CustomGrantAuthenticationToken.class.isAssignableFrom(authentication);
    }

    public void setTokenGenerator(OAuth2TokenGenerator<?> tokenGenerator) {
        Assert.notNull(tokenGenerator, "tokenGenerator cannot be null");
        this.tokenGenerator = tokenGenerator;
    }

    public void setAuthenticationManager(AuthenticationManager authenticationManager) {
        Assert.notNull(authorizationService, "authenticationManager cannot be null");
        this.authenticationManager = authenticationManager;
    }

    public void setAuthorizationService(OAuth2AuthorizationService authorizationService) {
        Assert.notNull(authorizationService, "authorizationService cannot be null");
        this.authorizationService = authorizationService;
    }
}
